#!/usr/bin/env python3
import os,sys,shutil,zarr,argparse
from pathlib import Path
def process_zarr_file(input_zarr_path,output_zarr_path):
    print(f"Processing zarr file: {input_zarr_path}")
    print(f"Output to: {output_zarr_path}")
    assert os.path.exists(input_zarr_path),f"Input zarr file not found: {input_zarr_path}"
    if os.path.exists(output_zarr_path):
        print(f"Removing existing output directory: {output_zarr_path}")
        shutil.rmtree(output_zarr_path)
    input_zarr=zarr.open(input_zarr_path,mode='r')
    output_zarr=zarr.open(output_zarr_path,mode='w')
    fields_to_keep=['ee/action','ee/state','qpos/action','qpos/state','eelocal/state','meta/episode_ends']
    print("Copying specified fields...")
    for field_path in fields_to_keep:
        if field_path not in input_zarr:
            print(f"  Warning: Field {field_path} not found in input file, skipping")
            continue
        print(f"  Copying field: {field_path}")
        field_data=input_zarr[field_path]
        output_path=field_path
        parent_path='/'.join(output_path.split('/')[:-1])
        if parent_path:
            if parent_path not in output_zarr:
                output_zarr.create_group(parent_path)
        if hasattr(field_data,'shape'):
            output_zarr.create_dataset(output_path,data=field_data[:],chunks=field_data.chunks,dtype=field_data.dtype,compressor=field_data.compressor,fill_value=field_data.fill_value if hasattr(field_data,'fill_value') else None)
        else:
            output_zarr.create_group(output_path)
            copy_group_recursive(field_data,output_zarr[output_path])
    if '.zgroup' in input_zarr:
        print("  Copying root .zgroup file")
        group_info=input_zarr.attrs.asdict()
        output_zarr.attrs.update(group_info)
    print("Field copying completed!")
    print("\nOutput zarr file structure:")
    print_zarr_structure(output_zarr)
    print(f"\nProcessing completed! New file saved at: {output_zarr_path}")
def copy_group_recursive(source_group,target_group):
    for key,item in source_group.items():
        if hasattr(item,'shape'):
            target_group.create_dataset(key,data=item[:],chunks=item.chunks,dtype=item.dtype,compressor=item.compressor,fill_value=item.fill_value if hasattr(item,'fill_value') else None)
        else:
            new_group=target_group.create_group(key)
            copy_group_recursive(item,new_group)
def print_zarr_structure(zarr_obj,indent=0):
    for key,item in zarr_obj.items():
        print("  "*indent+f"├── {key}")
        if hasattr(item,'shape'):
            print("  "*(indent+1)+f"    (array: {item.shape}, type: {item.dtype})")
        else:
            print_zarr_structure(item,indent+1)
def main():
    parser=argparse.ArgumentParser(description="Process zarr file, keeping only trajectory-related fields",formatter_class=argparse.RawDescriptionHelpFormatter,epilog="Example usage:\n  python zarr_for_traj.py /path/to/input.zarr\n  python zarr_for_traj.py /path/to/input.zarr /path/to/output.zarr")
    parser.add_argument('input_zarr',help='Input zarr file path')
    parser.add_argument('output_zarr',nargs='?',help='Output zarr file path (optional, default adds -traj to input filename)')
    args=parser.parse_args()
    input_path=Path(args.input_zarr).resolve()
    if args.output_zarr is None:
        output_path=input_path.parent/f"{input_path.stem}-traj.zarr"
    else:
        output_path=Path(args.output_zarr).resolve()
    try:
        process_zarr_file(str(input_path),str(output_path))
    except Exception as error:
        print(f"Error: {error}")
        import traceback
        traceback.print_exc()
        sys.exit(1)
if __name__=="__main__":
    main() 